# _*_ codeing : utf-8 _*_
# @Time : 2024-11-10
# @Author : wxp
# @File : __init__.py
# @Project : wxp-python

import ollama

from ollama import Client


# 模型客户端
def model_client(host):
    client = Client(host=host)
    return client


# 模型列表
def modelList(host):
    return model_client(host).list()


# 聊天接口
def chat(host, model, messages, stream):
    response = model_client(host).chat(
        model=model,
        messages=[
            {
                'role': 'user',
                'content': messages,
            },
        ],
        stream=stream,
    )
    print(response['message']['content'])
    return response


if __name__ == '__main__':
    model_list = modelList('http://localhost:11434')
    print("模型列表：", model_list)
    chat_result = chat('http://localhost:11434', 'llama3.2:1b', '为什么天空是蓝色的？', False)
    print("聊天响应: ", chat_result)
